#!/usr/bin/python3
# -*- coding:utf-8 -*-
import torch
import numpy as np
from einops import repeat
import os


def gen_golden_data_simple():
    # Parse tiling data in src/main.cpp
    """
    Code looks like:

    opDesc.num_rows = 18;
    opDesc.num_columns = 10;
    opDesc.batchShape = {3, 4};
    opDesc.dtype = 0;
    """
    num_rows = 0
    num_columns = 0
    batch_shape = []
    dtype = np.float32
    with open("./src/main.cpp", "r") as f:
        lines = f.readlines()
        for line in lines:
            if "opDesc.num_rows = " in line:
                num_rows = int(line.split("=")[1].strip().replace(";", ""))
            elif "opDesc.num_columns = " in line:
                num_columns = int(line.split("=")[1].strip().replace(";", ""))
            elif "opDesc.batchShape = " in line:
                batch_shape = list(
                    map(int, line.split("{")[1].split("}")[0].split(","))
                )
            elif "opDesc.dtype = " in line:
                dtype = int(line.split("=")[1].strip().replace(";", ""))
                if dtype == 0:
                    dtype = np.float32
                elif dtype == 1:
                    dtype = np.float16
                elif dtype == 2:
                    dtype = np.int32
                elif dtype == 3:
                    dtype = np.double

    print(
        f"num_rows: {num_rows}, num_columns: {num_columns}, batch_shape: {batch_shape}, dtype: {dtype}"
    )

    input_x = torch.randn(*batch_shape, num_rows, num_columns).numpy().astype(dtype)
    golden = torch.eye(
        n=num_rows,
        m=num_columns,
    )
    golden = (
        repeat(golden, "n m -> b n m", b=torch.prod(torch.tensor(batch_shape)))
        .numpy()
        .astype(dtype)
    )

    # Make directory
    if not os.path.exists("./input"):
        os.makedirs("./input")
    if not os.path.exists("./output"):
        os.makedirs("./output")

    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")


if __name__ == "__main__":
    gen_golden_data_simple()
